package com.timeriver.param_tuning

import org.apache.spark.ml.Pipeline
import org.apache.spark.ml.classification.LogisticRegression
import org.apache.spark.ml.evaluation.BinaryClassificationEvaluator
import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.ml.param.ParamMap
import org.apache.spark.ml.tuning.{ParamGridBuilder, TrainValidationSplit, TrainValidationSplitModel}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

/**
  * 简单交叉验证：结果带有随机性
  */
object TrainValidationSplitDemo {
  def main(args: Array[String]): Unit = {
    val session: SparkSession = SparkSession.builder()
      .appName("简单交叉验证")
      .master("local[6]")
      .getOrCreate()

    val df: DataFrame = session.read
      .format("jdbc")
      .option("url", "jdbc:mysql://10.0.24.197:3306/ml_datasets")
      .option("driver", "com.mysql.jdbc.Driver")
      .option("dbtable", "breast_cancer_wisconsin")
      .option("user", "root")
      .option("password", "123456")
      .load()

    /** 过滤缺失值 */
    val value: Dataset[Row] = df.filter(!_.anyNull)

    /** 获取特征列字段数组 */
    val inputCols: Array[String] = "clump_thickness,uniformity_of_cell_size,uniformity_of_cell_shape,marginal_adhesion,single_epithelial_cell_size,bare_nuclei,blan_chromatin,normal_nucleoli,mitoses".split(",")

    /** 构建特征列向量 */
    val data: DataFrame = new VectorAssembler()
      .setInputCols(inputCols)
      .setOutputCol("features")
      .transform(value)

    /** 注意：：：必须把标签列的列名改为label，否则在模型选择时会报label不存在 */
    val frame: DataFrame = data.withColumnRenamed("class", "label")

    /** 数据分割 */
    val Array(train, test) = frame.randomSplit(Array(0.9, 0.1), seed = 123)

    /** 创建逻辑回归模型 */
    val regression: LogisticRegression = new LogisticRegression()
      .setMaxIter(15)
      .setFeaturesCol("features")

    val pipeline = new Pipeline()
      .setStages(Array(regression))


    /** 构建参数网格搜索：参数调优 */
    val paramMaps: Array[ParamMap] = new ParamGridBuilder()
      .addGrid(regression.regParam, Array(0.1, 0.01))
      .addGrid(regression.fitIntercept)
      .addGrid(regression.elasticNetParam, Array(0.0, 0.5, 1.0))
      .build()

    /** 进行简单交叉验证 */
    val split: TrainValidationSplit = new TrainValidationSplit()
      .setEstimator(regression)
      .setEvaluator(new BinaryClassificationEvaluator())
      .setEstimatorParamMaps(paramMaps)
      /** 80%的数据将用于训练，其余20%用于验证 */
      .setTrainRatio(0.8)
      .setParallelism(2)

    val model: TrainValidationSplitModel = split.fit(train)

    /** 模型保存，可以直接将整个pipeline管道存储起来 */
    model.save("./model/logisticByTrainValidation")

    val res: DataFrame = model.transform(test)
    res.show(5,false)

    /** 二分类模型评估 */
    val evaluator: BinaryClassificationEvaluator = new BinaryClassificationEvaluator()
      .setLabelCol("label")
      .setRawPredictionCol("rawPrediction")
      .setMetricName("areaUnderROC")

    val d: Double = evaluator.evaluate(res)
    println("areaUnderROC=", d)
  }
}
